# -*- coding: utf-8 -*-


from locust import HttpLocust, TaskSet, task
import random
import os
from faker import Faker
from locust import HttpUser, task, between, HttpLocust
from locust import LoadTestShape
import requests
from requests.packages.urllib3.exceptions import InsecureRequestWarning

requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
fake = Faker(locale='zh_CN')
import json


class QuickstartUser(HttpUser):
    host = " http://8.136.194.239"

    def on_start(self):
        pass

    @task
    def api1(self):
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.163 Safari/537.36',
            'Content-Type': 'application/json;charset=UTF-8',
        }

        with self.client.get(url='/dev/short-url/su/query/9BNUwc', headers=headers, catch_response=True) as response:
            if response.status_code == 200:
                try:
                    assert json.loads(response.text)['errCode'] == 0
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)
                response.failure("No data")

    # @task
    def api2(self):
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.163 Safari/537.36',
            'Content-Type': 'application/json;charset=UTF-8',
        }
        payload = {
            "originalUrl": "http://user-dev.tangees.com/infinite-group-qrcode&taskId=1463792629605257217"
        }

        with self.client.post(url='/dev/short-url/su/create', headers=headers, data=json.dumps(payload), catch_response=True) as response:
            if response.status_code == 200:
                print(response.text)

                try:
                    json.loads(response.text)['errCode'] == 0
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)
                # response.failure("No data")

    def on_stop(self):
        pass


class StagesShape(LoadTestShape):
    """
    在不同的阶段 具有不同的用户数和 产生率的 图形形状


    参数解析:
        stages ：字典列表，每个字典都具有下列这些键值的阶段:
            duration -- 持续时间，  经过多少秒后，进入到下个阶段
            users -- 总用户数
            spawn_rate -- 产生率，即每秒启动/停止的用户数
            stop -- 可以在特定阶段停止测试的值
        stop_at_end -- 可以在所有阶段设置运行后停止

    """

    stages = [
        {"duration": 5, "users": 20, "spawn_rate": 1},
        # {"duration": 300, "users": 10, "spawn_rate": 2},
        # {"duration": 300, "users": 20, "spawn_rate": 2}
    ]

    def tick(self):

        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage['duration']:
                tick_data = (stage['users'], stage['spawn_rate'])
                return tick_data

        return None


if __name__ == "__main__":
    os.system("locust -f locust_短链.py --web-host=127.0.0.1 --web-port=7717")
